from typing import *
import sympy
from sympy.plotting import plot_implicit
from sympy.parsing.sympy_parser import parse_expr

###############################################################
########################函数定义区##############################
###############################################################
def get_single_line(expression:str,color:str):
    """get_single_line用于绘制单条曲线"""
    lhs, rhs = expression.replace("^","**").split("=")
    eqn_lhs = parse_expr(lhs)
    eqn_rhs = parse_expr(rhs)
    return plot_implicit(eqn_lhs-eqn_rhs,line_color=color,show=False)

def draw_flow_field(streamlines, color="#eb5523"): 
    """
    说明：
        draw_flow_field用于可视化流场中的曲线，包括流线、迹线和脉线。
    参数：
        streamlines：str或list，需要可视化的曲线
        color： 绘制的曲线的颜色，用16进制表示
    返回值：
        plot：绘制得到的图像，可以被用于后续的保存等操作
    """
    assert type(streamlines) in [str,list], f"streamline参数要求实参为str或者list，接受到的实参类型为{type(streamlines)}"
    if type(streamlines) == str:
        lines = [get_single_line(streamlines,color)]
    else:
        lines = [get_single_line(line,color) for line in streamlines]
    plot = lines[0]
    for i in range(1,len(lines)):
        plot.append(lines[i][0])
    return plot

def save_fig(plot,filename):
    """
    说明：
        save_fig用于保存绘制得到的图像
    参数：
        plot：需要保存的图像，必须是draw_flow_field的返回值
        filename：保存的图像的名字
    """
    assert type(filename) == str,"必须以字符串形式指定保存的文件名"
    plot.save(filename)

def show_fig(plot):
    plot.show()



###############################################################
############################主程序##############################
###############################################################
if __name__ == "__main__":
    # 绘制曲线过多时耗时较长，需等待10秒左右
    streamlines = ['20*y*(1-16/(x^2+y^2))+0.3*exp(sqrt(x^2+y^2))/(2*pi)='+str(i) for i in range(-30,31,3)]
    p = draw_flow_field(streamlines, color="#20a482")
    show_fig(p)
    save_fig(p,"有环量圆柱扰流.png")